#include <string.h>
#include "Package.h"
#include "MqttValidor.h"
#include "Core/Logger.h"
#include "MQTT/MQTTError.h"

static int32_t parsePropertyElement(MQTTEntity *entity, uint8_t *property, uint32_t *idx,
                                    ValueLength *tlv, UserProperty *prop, uint32_t *integer);


int32_t mqttConnectPkt(MQTTEntity *entity, uint8_t **pkt, uint32_t *pktLen, ConnectProperty *prop) {

    uint32_t idx = 0, connPaylaodLen = 0, connRemainlen = 0, connPktLen = 0, propertyLen = 0, propertyTotalLen = 0;
    uint8_t *pos = nullptr;
    const uint8_t connFixedHeader = MQTT_CONN_PKT_TYPE;
    const uint8_t connProtocolName[] = {0x00, 0x04, 0x4D, 0x51, 0x54, 0x54};
    uint8_t connProtocolLevel = 0x04;
    const uint8_t connConnectFlag = 0xC0 | (entity->cleanSession << 1);
    uint32_t propertyLenOffset = 0;
    uint8_t propertyLenArray[4] = {0};
    GeneralProperty generalProperty = {0};

    /** property len*/
    connProtocolLevel = 0x5;
    if (nullptr != prop) {
        generalProperty.maxTopicAlias = prop->topicAliasMax;
        memcpy(&(generalProperty.userProperty[0]), &(prop->userProperty[0]),MAX_USER_PROPERTY * (sizeof(UserProperty*)));
    }

    if (nullptr == entity->preConnectProperty) {
        /* Normal Connection Case */
        int res = mqttGeneralPropertyValid(&generalProperty);
        if (res != aiotSuccess) {
            return res;
        }

        propertyLen = getGeneralPropertyLen(&generalProperty);
        encodeRemainLen(propertyLen, propertyLenArray, &propertyLenOffset);
    } else {
        /** Reconnetion Case: Use Connection Properties Stored Previously */
        uint8_t tmp_offset = 0;
        readVariableByteInterger(entity->preConnectProperty, &propertyLen, &tmp_offset);
        propertyLenOffset = tmp_offset;
    }

    propertyTotalLen = propertyLenOffset + propertyLen;
    /** payload length */
    connPaylaodLen = (uint32_t)(strlen(entity->clientId) + strlen(entity->username) + strlen(entity->password) + 3 * MQTT_UTF8_STR_EXTRA_LEN);

    /** remain length value */
    connRemainlen = MQTT_CONN_REMAINLEN_FIXED_LEN + connPaylaodLen + propertyTotalLen;

    /** Total Packet Length */
    connPktLen = MQTT_CONN_FIXED_HEADER_TOTAL_LEN + connPaylaodLen + propertyTotalLen;

    pos = getSystemInterface()->memMalloc(connPktLen);
    if (pos == nullptr) {
        return memoryMallocFailed;
    }

    memset(pos, 0, connPktLen);

    /* Fixed Header */
    pos[idx++] = connFixedHeader;

    /** remain length */
    encodeRemainLen(connRemainlen, pos + idx, &idx);

    /** protocol name */
    memcpy(pos + idx, connProtocolName, MQTT_CONN_PROTOCOL_NAME_LEN);
    idx += MQTT_CONN_PROTOCOL_NAME_LEN;

    /* Protocol Level */
    pos[idx++] = connProtocolLevel;

    /* Connect Flag */
    pos[idx++] = connConnectFlag;

    /* Keep Alive MSB */
    pos[idx++] = (uint8_t)((entity->aiotEntity.keepAlive >> 8) & 0x00FF);

    /* Keep Alive LSB */
    pos[idx++] = (uint8_t)((entity->aiotEntity.keepAlive) & 0x00FF);

    /* property */
    if (nullptr != entity->preConnectProperty) {
        /** normal connection case: use connection properties stored previously */
        memcpy(&(pos[idx]), entity->preConnectProperty, propertyTotalLen);
        idx += propertyTotalLen;
    } else  {
        /** not reconnection case */
        uint8_t *start = &(pos[idx]);
        writeGeneralProp(pos, &idx, &generalProperty, propertyLenArray, propertyLenOffset);
        entity->preConnectProperty = getSystemInterface()->memMalloc(propertyTotalLen);
        if (nullptr == entity->preConnectProperty) {
            return memoryMallocFailed;
        }

        memset(entity->preConnectProperty, 0, propertyTotalLen);
        memcpy(entity->preConnectProperty, start, propertyTotalLen);
    }

    /** payload: clientId, username, password */
    setUtf8EncodedString((uint8_t *) entity->clientId, strlen(entity->clientId), pos + idx);
    idx += MQTT_UTF8_STR_EXTRA_LEN + strlen(entity->clientId);

    setUtf8EncodedString((uint8_t *) entity->username, strlen(entity->username), pos + idx);
    idx += MQTT_UTF8_STR_EXTRA_LEN + strlen(entity->username);

    setUtf8EncodedString((uint8_t *) entity->password, strlen(entity->password), pos + idx);
    idx += MQTT_UTF8_STR_EXTRA_LEN + strlen(entity->password);

    *pkt = pos;
    *pktLen = idx;

    return aiotSuccess;
}

void writeGeneralProp(uint8_t *pos, uint32_t *index, GeneralProperty *generalProperty,uint8_t *propertyLen, uint8_t propertyLenOffset) {

    uint32_t idx = *index;
    uint16_t topicAliasMax = generalProperty->maxTopicAlias;
    uint16_t topicAlias = generalProperty->topicAlias;
    uint16_t clientReceiveMax = generalProperty->maxClientReceive;
    MQTTValue *reasonString = generalProperty->reasonString;
    UserProperty **propList = generalProperty->userProperty;

    /** Total Length of All Properties */
    if (nullptr != propertyLen) {
        memcpy(&(pos[idx]), propertyLen, propertyLenOffset);
        idx += propertyLenOffset;
    }

    /** User Properties */
    if (nullptr != propList) {
        int iter = 0;
        for (iter = 0; iter < MAX_USER_PROPERTY; iter++) {
            UserProperty *prop = propList[iter];
            if (nullptr != prop) {
                pos[idx++] = MQTTPROP_USER_PROPERTY;
                pos[idx++] = prop->key.len << 8;
                pos[idx++] = prop->key.len & 0xff;
                memcpy(&(pos[idx]), prop->key.value, prop->key.len);
                idx += prop->key.len;

                pos[idx++] = prop->value.len << 8;
                pos[idx++] = prop->value.len & 0xff;
                memcpy(&(pos[idx]), prop->value.value, prop->value.len);
                idx += prop->value.len;
            }
        }
    }

    /** User Topic Alias max */
    if (topicAliasMax > 0) {
        pos[idx++] = MQTTPROP_TOPIC_ALIAS_MAX;
        pos[idx++] = topicAliasMax >> 8;
        pos[idx++] = topicAliasMax & 0xFF;
    }

    /** Topic Alias */
    if (topicAlias > 0) {
        pos[idx++] = MQTTPROP_TOPIC_ALIAS;
        pos[idx++] = topicAlias >> 8;
        pos[idx++] = topicAlias & 0xFF;
    }

    /** Client Receive Max */
    if (clientReceiveMax > 0) {
        pos[idx++] = MQTTPROP_RECEIVE_MAXIMUM;
        pos[idx++] = clientReceiveMax >> 8;
        pos[idx++] = clientReceiveMax & 0xFF;
    }

    /** Response Topic */
    if (0 != generalProperty->responseTopic.len) {
        pos[idx++] = MQTTPROP_RESPONSE_TOPIC;
        uint16_t topicLen = generalProperty->responseTopic.len;
        pos[idx++] = topicLen >> 8;
        pos[idx++] = topicLen & 0xFF;
        memcpy(&(pos[idx]), generalProperty->responseTopic.value, topicLen);
        idx += topicLen;
    }

    /** Correlation Data */
    if (0 != generalProperty->correlationData.len) {
        pos[idx++] = MQTTPROP_CORRELATION_DATA;
        uint16_t dataLen = generalProperty->correlationData.len;
        pos[idx++] = dataLen >> 8;
        pos[idx++] = dataLen & 0xFF;
        memcpy(&(pos[idx]), generalProperty->correlationData.value, dataLen);
        idx += dataLen;
    }

    /** Subscription Identifier */
    if (generalProperty->subIdentifier > 0) {
        uint32_t subIdentiferOffset = 0;
        uint8_t  subIdentifier[4] = {0};
        encodeRemainLen(generalProperty->subIdentifier, &subIdentifier[0],&subIdentiferOffset);
        pos[idx++] = MQTTPROP_SUBSCRIPTION_IDENTIFIER;
        memcpy(&pos[idx], subIdentifier, subIdentiferOffset);
        idx += subIdentiferOffset;
    }

    /** Reason String */
    if (nullptr != reasonString) {
        pos[idx++] =  MQTTPROP_REASON_STRING;
        pos[idx++] = (reasonString->len) >> 8 ;
        pos[idx++] =  reasonString->len & 0xFF;
        memcpy(&pos[idx], reasonString->value, reasonString->len);
        idx += reasonString->len ;
    }

    *index = idx;
}

void encodeRemainLen(uint32_t input, uint8_t *output, uint32_t *outputIdx) {

    uint8_t encoded_byte = 0, idx = 0;

    do {
        encoded_byte = input % 128;
        input /= 128;
        if (input > 0) {
            encoded_byte |= 128;
        }

        output[idx++] = encoded_byte;
    } while (input > 0);

    *outputIdx += idx;
}

void setUtf8EncodedString(uint8_t *input, uint16_t inputLen, uint8_t *output) {

    uint32_t idx = 0, input_idx = 0;

    /** String Length MSB */
    output[idx++] = (uint8_t)((inputLen >> 8) & 0x00FF);

    /** String Length LSB */
    output[idx++] = (uint8_t)((inputLen) & 0x00FF);

    /** UTF-8 Encoded Character Data */
    for (input_idx = 0; input_idx < inputLen; input_idx++) {
        output[idx++] = input[input_idx];
    }
}

int32_t readVariableByteInterger(uint8_t *input, uint32_t *remainlen, uint8_t *offset) {

    uint8_t ch = 0;
    uint32_t multiplier = 1;
    uint32_t mqtt_remainlen = 0;
    uint8_t pos = 0;

    do {
        ch = input[pos++];
        mqtt_remainlen += (ch & 127) * multiplier;
        if (multiplier > 128 * 128 * 128) {
            return mqttMalformedRemainingLen;
        }

        multiplier *= 128;
    } while ((ch & 128) != 0 && pos < 4);

    *remainlen = mqtt_remainlen;
    *offset = pos;

    return aiotSuccess;
}

int32_t parsePublishProperties(MQTTEntity *entity, uint8_t *property,
                               PubProperty *prop, uint32_t *propertyLen, uint8_t *parsedOffset, uint32_t remainLen) {

    GeneralProperty general_prop = {0};
    int32_t res = parseGeneralProperties(entity, property, &general_prop, propertyLen,
                                         parsedOffset, remainLen);
    if (res != aiotSuccess) {
        return res;
    }

    prop->expireInterval = general_prop.msgExpireInterval;
    prop->topicAlias = general_prop.topicAlias;
    memcpy(&(prop->userProperty[0]), &(general_prop.userProperty[0]),MAX_USER_PROPERTY * (sizeof(UserProperty *)));
    prop->subIdentifier = general_prop.subIdentifier;
    prop->responseTopic = general_prop.responseTopic;
    prop->correlationData = general_prop.correlationData;
    return aiotSuccess;
}

int32_t parseConAckProperties(MQTTEntity *entity, uint8_t *property, ConnAckProperty *prop, uint32_t remainLen) {

    GeneralProperty generalProp = {0};
    int32_t res = parseGeneralProperties(entity, property, &generalProp, nullptr, nullptr, remainLen);
    if (res != aiotSuccess) {
        return res;
    }

    prop->maxQos = generalProp.maxQos;
    prop->maxTopicAlis = generalProp.maxTopicAlias;
    memcpy(&(prop->userProperty[0]), &(generalProp.userProperty[0]),MAX_USER_PROPERTY * (sizeof(UserProperty *)));
    prop->clientId = generalProp.assignedCliId;
    prop->maxPacket = generalProp.maxPacketSize;
    prop->serverReceiveMax = generalProp.serverReceiveMax;
    prop->wildcardSubAvailable   = generalProp.wildSubAvailable;
    prop->subIdentifierAvailable = generalProp.subIdentifierAvailable;
    prop->sharedSubAvailable = generalProp.sharedSubAvailable;
    entity->txTopicAliasSize = prop->maxTopicAlis;

    return aiotSuccess;
}


int32_t parseGeneralProperties(MQTTEntity *mqtt_handle, uint8_t *property,
                               GeneralProperty *prop, uint32_t *propertyLen, uint8_t *offset, uint32_t len) {

    uint32_t idx = 0, total_prop_len = 0;
    int32_t res;
    UserProperty **user_prop = prop->userProperty;
    uint32_t property_len = 0;
    uint8_t property_len_offset = 0;
    ValueLength tlv = {0};
    UserProperty cur_user_prop;
    uint8_t user_property_count = 0;
    uint32_t variable_byte_integer = 0;

    res = readVariableByteInterger(&(property[0]), &property_len, &property_len_offset);
    if (res != aiotSuccess) {
        return res;
    }

    /** quit if the total len of all properties exceeds the remain_len of mqtt packet */
    if (property_len > len - property_len_offset) {
        return mqttInvalidPropertyLen;
    }

    if (nullptr != propertyLen) {
        *propertyLen = property_len;
    }

    if (nullptr != offset) {
        *offset = property_len_offset;
    }

    total_prop_len = property_len + property_len_offset;
    if (property_len == 0) {
        return aiotSuccess;
    }

    idx += property_len_offset;
    prop->maxQos = 2;                            /* by default is 2 */
    prop->wildSubAvailable = 1;    /* enabled by default */
    prop->subIdentifierAvailable = 1;  /* enabled by default */
    prop->sharedSubAvailable = 1;      /* enabled by default */

    while (total_prop_len > idx) {

        memset(&tlv, 0, sizeof(tlv));
        memset(&cur_user_prop, 0, sizeof(cur_user_prop));
        res = parsePropertyElement(mqtt_handle, property + idx, &idx, &tlv, &cur_user_prop, &variable_byte_integer);
        if (res != aiotSuccess) {
            return res;
        }

        /** quit if the length of a certain property exceeds the remain_len of the whole packet */
        if (idx > len) {
            return mqttInvalidPropertyLen;
        }

        switch (tlv.type) {
            case MQTTPROP_TOPIC_ALIAS: {
                uint16_t topic_alias = 0;
                topic_alias = (uint16_t)(*(tlv.value) << 8) ;
                topic_alias += (uint16_t)(*(tlv.value + 1));
                prop->topicAlias = topic_alias;
            }
                break;
            case MQTTPROP_TOPIC_ALIAS_MAX: {
                uint16_t topic_alias_max = 0;
                topic_alias_max = (uint16_t)(*(tlv.value) << 8) ;
                topic_alias_max += (uint16_t)(*(tlv.value + 1));
                prop->maxTopicAlias = topic_alias_max;
            }
                break;
            case MQTTPROP_MAX_QOS: {
                uint8_t max_qos = 0;
                max_qos = (uint8_t)(*(tlv.value));
                prop->maxQos = max_qos;
            }
                break;
            case MQTTPROP_ASSIGNED_CLIENT_IDENTIFIER: {
                void *cid = getSystemInterface()->memMalloc(tlv.len + 1);
                if (nullptr == cid) {
                    return memoryMallocFailed;
                }

                memset(cid, 0, tlv.len + 1);
                memcpy(cid, tlv.value, tlv.len);
                prop->assignedCliId = cid;
            }
                break;
            case MQTTPROP_WILDCARD_SUBSCRIPTION_AVAILABLE: {
                prop->wildSubAvailable = (uint8_t)(*(tlv.value));
            }
                break;
            case MQTTPROP_SUBSCRIPTION_IDENTIFIER_AVAILABLE: {
                prop->subIdentifierAvailable = (uint8_t)(*(tlv.value));
            }
                break;
            case MQTTPROP_SHARED_SUBSCRIPTION_AVAILABLE: {
                prop->sharedSubAvailable = (uint8_t)(*(tlv.value));
            }
                break;
            case MQTTPROP_PUBLICATION_EXPIRY_INTERVAL: {
                uint32_t message_expire = 0;
                message_expire += *(tlv.value) << 24;
                message_expire += *(tlv.value + 1) << 16;
                message_expire += *(tlv.value + 2) << 8;
                message_expire += *(tlv.value + 3) << 0;
                prop->msgExpireInterval = message_expire;
            }
                break;
            case MQTTPROP_MAX_PACK_SIZE: {
                uint32_t max_pack_size = 0;
                max_pack_size += *(tlv.value) << 24;
                max_pack_size += *(tlv.value + 1) << 16;
                max_pack_size += *(tlv.value + 2) << 8;
                max_pack_size += *(tlv.value + 3) << 0;
                prop->maxPacketSize = max_pack_size;
            }
                break;
            case MQTTPROP_RECEIVE_MAXIMUM: {
                uint16_t server_receive_max = 0;
                server_receive_max = (uint16_t)(*(tlv.value) << 8) ;
                server_receive_max += (uint16_t)(*(tlv.value + 1));
                prop->serverReceiveMax = server_receive_max;
            }
                break;
            case MQTTPROP_USER_PROPERTY: {
                if (user_property_count < MAX_USER_PROPERTY && cur_user_prop.key.len != 0) {
                    user_prop[user_property_count] = getSystemInterface()->memMalloc(sizeof(UserProperty));
                    if (nullptr == user_prop[user_property_count]) {
                        return memoryMallocFailed;
                    }

                    memcpy(user_prop[user_property_count], &cur_user_prop, sizeof(UserProperty));
                    user_property_count++;
                }
            }
                break;
            case MQTTPROP_UNRESOLVED: {
                return mqttUnknownPropertyOption;
            }
            default:
                break;
        }
    }

    return aiotSuccess;
}

static int32_t parsePropertyElement(MQTTEntity *entity, uint8_t *property, uint32_t *idx,ValueLength *tlv, UserProperty *prop, uint32_t *integer) {

    uint32_t readLen = 0;
    uint8_t type = property[0];
    int32_t res;

    switch (type) {
        /* Bits (Single Byte) */
        case MQTTPROP_PAYLOAD_FORMAT_INDICATOR:
        case MQTTPROP_REQUEST_PROBLEM_INFORMATION:
        case MQTTPROP_REQUEST_RESPONSE_INFORMATION:
        case MQTTPROP_MAX_QOS:
        case MQTTPROP_RETAIN_AVAILABLE:
        case MQTTPROP_WILDCARD_SUBSCRIPTION_AVAILABLE:
        case MQTTPROP_SUBSCRIPTION_IDENTIFIER_AVAILABLE:
        case MQTTPROP_SHARED_SUBSCRIPTION_AVAILABLE: {
            readLen++;
            tlv->type = type;
            tlv->len = 1;
            readLen += 1;
            tlv->value = property + MQTT_V5_PROPERTY_ID_LEN ;
        }
            break;

            /* Two Byte Integer */
        case MQTTPROP_SERVER_KEEP_ALIVE:
        case MQTTPROP_TOPIC_ALIAS:
        case MQTTPROP_RECEIVE_MAXIMUM:
        case MQTTPROP_TOPIC_ALIAS_MAX: {
            readLen++;
            tlv->type = type;
            tlv->len = 2;
            readLen += 2;
            tlv->value = property + MQTT_V5_PROPERTY_ID_LEN;
        }
            break;

            /* Four Byte Integer */
        case MQTTPROP_PUBLICATION_EXPIRY_INTERVAL:
        case MQTTPROP_SESSION_EXPIRY_INTERVAL:
        case MQTTPROP_WILL_DELAY_INTERVAL:
        case MQTTPROP_MAX_PACK_SIZE: {
            readLen++;
            tlv->type = type;
            tlv->len = 4;
            readLen += 4;
            tlv->value = property + MQTT_V5_PROPERTY_ID_LEN;
        }
            break;

            /* Variable Byte Integer */
        case MQTTPROP_SUBSCRIPTION_IDENTIFIER: {
            readLen++;
            uint8_t *remainLenStart = property + readLen;
            uint32_t remainLen = 0;
            uint8_t offset = 0;

            res = readVariableByteInterger(remainLenStart, &remainLen, &offset);
            if (res != aiotSuccess) {
                return res;
            }
            readLen += offset;

            tlv->type = type;
            tlv->len = offset;
            *integer = remainLen;
        }
            break;

            /* Binary Data; UTF-8 String */
        case MQTTPROP_CONTENT_TYPE:
        case MQTTPROP_RESPONSE_TOPIC:
        case MQTTPROP_ASSIGNED_CLIENT_IDENTIFIER:
        case MQTTPROP_AUTHENTICATION_METHOD:
        case MQTTPROP_RESPONSE_INFORMATION:
        case MQTTPROP_SERVER_REFERENCE:
        case MQTTPROP_REASON_STRING:
        case MQTTPROP_CORRELATION_DATA:
        case MQTTPROP_AUTHENTICATION_DATA: {
            readLen++;
            tlv->type = type;
            tlv->len = ((*(property + 1)) << 8) + (*(property + 2));
            readLen += 2;
            tlv->value = property + 3;
            readLen += tlv->len;
        }
            break;

            /* UTF-8 String Pairs. especially for User Properties */
        case MQTTPROP_USER_PROPERTY: {
            readLen++;
            uint8_t *keyStart = property + readLen;
            prop->key.len = ((*(keyStart)) << 8) + (*(keyStart + 1));
            prop->key.value = keyStart + 2;
            readLen += 2 + prop->key.len;

            uint8_t *valueStart = property + readLen;
            prop->value.len = ((*(valueStart)) << 8) + (*(valueStart + 1));
            prop->value.value = valueStart + 2;
            readLen += 2 + prop->value.len;

            tlv->type = type;
            tlv->len = 2 + prop->key.len + 2 + prop->value.len;
        }
            break;

        default: {
            tlv->type = MQTTPROP_UNRESOLVED;
            loggerOut(aiotDebug,__FILE__,__LINE__,"unresolved property option\n");
        }
            break;
    }

    *idx = *idx + readLen;
    return aiotSuccess;
}
